//! knn的实例代码
//! datingTestSet.csv文件是样本数据
//!
use csv::ReaderBuilder;
use linfa::prelude::*;
use linfa_preprocessing::linear_scaling::*;
use machine_study::{
    knn::{BallTree, Knn, L1Dist, LinearSearch},
    split_data_to_train_and_test,
};
use ndarray::prelude::*;
use ndarray_csv::Array2Reader;
use plotters::prelude::*;
fn main() {
    //test_knn_plotters();
    //test_knn_simple();
    test_knn_csv();
}

fn test_knn_simple() {
    let features_data = array![
        [3.393533211, 2.331273381],
        [3.110073483, 1.781539638],
        [1.343808831, 3.368360954],
        [3.582294042, 4.679179110],
        [2.280362439, 2.866990263],
        [7.423436942, 4.696522875],
        [5.745051997, 3.533989803],
        [9.172168622, 2.511101045],
        [7.792783481, 3.424088941],
        [7.939820817, 0.791637231],
    ];

    let target_data = array![0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0];
    let lp = LinearScalerParams::new(ScalingMethod::Standard(true, true));
    let data_set = DatasetBase::new(features_data.view(), target_data.view());
    let lp = lp.fit(&data_set).unwrap();
    let ret = lp.transform(data_set);
    println!("ret:\n{}", ret.records());
    println!("{}", features_data);
    println!("{}", target_data);
    let new_data = array![[8.093607318, 3.365731514]];
    let knn = Knn::fit(&features_data, &target_data, BallTree);
    let p = knn.predict(&new_data, 10, L1Dist).unwrap();
    println!("p:{}", p);
}

fn test_knn_plotters() {
    let features_data = array![
        [3.393533211, 2.331273381],
        [3.110073483, 1.781539638],
        [1.343808831, 3.368360954],
        [3.582294042, 4.679179110],
        [2.280362439, 2.866990263],
        [7.423436942, 4.696522875],
        [5.745051997, 3.533989803],
        [9.172168622, 2.511101045],
        [7.792783481, 3.424088941],
        [7.939820817, 0.791637231],
    ];

    let target_data = array![0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0];

    let root = SVGBackend::new("examples/knn/dating.svg", (640, 480)).into_drawing_area();
    root.fill(&WHITE).unwrap();

    let mut chart = ChartBuilder::on(&root)
        .x_label_area_size(40)
        .y_label_area_size(40)
        .caption("数据分布", ("Scan", 20))
        .build_cartesian_2d(0.0f64..10.0, 0.0f64..10.0)
        .unwrap();
    chart
        .configure_mesh()
        .disable_x_mesh()
        .disable_y_mesh()
        .draw()
        .unwrap();
    chart
        .draw_series(
            features_data
                .axis_iter(Axis(0))
                .zip(target_data.iter())
                .map(|v| {
                    Circle::new((v.0[0], v.0[1]), 5, {
                        if *v.1 == 0.0 {
                            RED.filled()
                        } else {
                            BLUE.filled()
                        }
                    })
                }),
        )
        .unwrap();
    let _new_data = array![[8.093607318, 3.365731514]];
    chart
        .draw_series(
            (0..1)
                .into_iter()
                .map(|_| Circle::new((8.093607318, 3.365731514), 5, GREEN.filled())),
        )
        .unwrap();
}

fn test_knn_csv() {
    let mut csv_reader = ReaderBuilder::default()
        .has_headers(false)
        .delimiter(b',')
        .from_path("./examples/knn/datingTestSet.csv")
        .unwrap();
    let data: Array2<f64> = csv_reader.deserialize_array2_dynamic().unwrap();
    let train_and_test = split_data_to_train_and_test(&data, 0.80).unwrap();
    // 数据归一化
    let lcp = LinearScalerParams::new(ScalingMethod::Standard(true, true));
    let ls = lcp
        .fit(&DatasetBase::new(
            train_and_test.x_train().view(),
            train_and_test.y_train(),
        ))
        .unwrap();
    let x_train = ls.transform(train_and_test.x_train().clone());
    let y_train = train_and_test.y_train();

    let x_test = ls.transform(train_and_test.x_test().clone());
    let y_test = train_and_test.y_test();
    let y_test_usize = y_test.map(|x| *x as usize);

    let test = DatasetBase::new(x_test.view(), y_test_usize.view())
        .with_feature_names(vec!["test1", "test2", "test3"]);

    // 寻找合适的K值
    let knn = Knn::fit(&x_train, &y_train, LinearSearch);
    let mut score = 0.0;
    let mut current_k = 0;
    for k in 1..5 {
        let predicts = knn
            .predict(&x_test, k, L1Dist)
            .unwrap()
            .map(|x| *x as usize);

        let confusion = predicts.confusion_matrix(&test).unwrap();
        if confusion.accuracy() > score {
            score = confusion.accuracy();
            current_k = k;
        }

        println!("score: {}, current_k: {}", confusion.accuracy(), k);
    }
    println!("最佳的k值:{}, 最佳的准确率: {}", current_k, score);
}
